from asyncio import FastChildWatcher
from code import interact
import os
import time
import numpy as np
from isaacgym import gymutil
from isaacgym import gymapi
from isaacgym import gymtorch
from math import sqrt
import math
from sympy import false
import torch
import cv2
import threading

from draw import *

from pcgworker.PCGWorker import *

from wfc_env import *

from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.results_plotter import load_results, ts2xy
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common import results_plotter
from stable_baselines3.common.results_plotter import load_results, ts2xy, plot_results
from stable_baselines3.common.noise import NormalActionNoise
from stable_baselines3.common.callbacks import BaseCallback

# define a thread which takes input
class InputThread(threading.Thread):
    def __init__(self):
        super(InputThread, self).__init__()
        self.daemon = True
        self.last_user_input = None

    def run(self):
        while True:
            self.last_user_input = input('input something: ')
            # do something based on the user input here
            print("self.last_user_input : ",self.last_user_input)

class SaveOnBestTrainingRewardCallback(BaseCallback):
    """
    Callback for saving a model (the check is done every ``check_freq`` steps)
    based on the training reward (in practice, we recommend using ``EvalCallback``).

    :param check_freq:
    :param log_dir: Path to the folder where the model will be saved.
      It must contains the file created by the ``Monitor`` wrapper.
    :param verbose: Verbosity level.
    """
    def __init__(self, check_freq: int, log_dir: str, verbose: int = 1):
        super(SaveOnBestTrainingRewardCallback, self).__init__(verbose)
        self.check_freq = check_freq
        self.log_dir = log_dir
        self.save_path = os.path.join(log_dir, 'best_model')
        self.best_mean_reward = -np.inf

        # start timer
        self.start = time.time()

    def _init_callback(self) -> None:
        # Create folder if needed
        if self.save_path is not None:
            os.makedirs(self.save_path, exist_ok=True)
        self.start = time.time()

    def _on_step(self) -> bool:
        if self.n_calls % self.check_freq == 0:
            end = time.time()
            
            # Retrieve training reward
            x, y = ts2xy(load_results(self.log_dir), 'timesteps')
            if len(x) > 0:
                # Mean training reward over the last 100 episodes
                mean_reward = np.mean(y[-100:])
                if self.verbose > 0:
                    print("Time taken: {}".format(end - self.start))
                    print(f"Num timesteps: {self.num_timesteps}")
                    print(f"Best mean reward: {self.best_mean_reward:.2f} - Last mean reward per episode: {mean_reward:.2f}")

                # New best model, you could save the agent here
                if mean_reward > self.best_mean_reward:
                    self.best_mean_reward = mean_reward
                    # Example for saving best model
                    if self.verbose > 0:
                        print(f"Saving new best model to {self.save_path}")
                        self.model.save(self.save_path)

        return True

it = InputThread()
it.start()

LOGDIR = "./training_logs"
timesteps = 100000

# Create the callback: check every 1000 steps
callback = SaveOnBestTrainingRewardCallback(check_freq=1000, log_dir=LOGDIR)

m_env = Monitor(CustomEnv(), LOGDIR, allow_early_resets=True)
observation = m_env.reset()

policy_kwargs = dict(activation_fn=torch.nn.ReLU,
                     net_arch=[dict(pi=[256, 256], vf=[256, 256])])

# model = PPO('CnnPolicy', m_env, policy_kwargs=policy_kwargs).learn(total_timesteps=timesteps,callback=callback)
model = PPO('CnnPolicy', m_env).learn(total_timesteps=timesteps,callback=callback)

plot_results([LOGDIR], timesteps, results_plotter.X_TIMESTEPS, "PPO in Isaac Gym")
plt.show()
